package cn.smileyan.demos.io;

import lombok.extern.slf4j.Slf4j;
import org.apache.flink.api.java.utils.MultipleParameterTool;
import org.apache.flink.connector.base.DeliveryGuarantee;
import org.apache.flink.connector.kafka.sink.KafkaRecordSerializationSchema;
import org.apache.flink.connector.kafka.sink.KafkaSink;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.kafka.clients.admin.AdminClient;
import org.apache.kafka.clients.admin.AdminClientConfig;
import org.apache.kafka.clients.admin.NewTopic;
import org.apache.kafka.clients.producer.ProducerConfig;

import java.util.Collections;
import java.util.Properties;
import java.util.concurrent.ExecutionException;

/**
 * 通过参数构建通用的 kafka 通讯序列化与反序列化实体
 * @author smileyan
 */
@Slf4j
public class CommonKafkaBuilder {
    /**
     * 构建参数
     */
    private final MultipleParameterTool parameterTool;

    public CommonKafkaBuilder(final MultipleParameterTool parameterTool) {
        this.parameterTool = parameterTool;
    }

    /**
     * 构建kafka sink
     * @param clazz 实体类class
     * @param <E> 实体类泛型
     * @return kafka sink 对象
     */
    public <E> KafkaSink<E> buildSink(Class<E> clazz) {
        final String bs = parameterTool.getRequired(KafkaArgs.BOOTSTRAP_SERVER.key);
        final String ot = parameterTool.getRequired(KafkaArgs.OUTPUT_TOPIC.key);
        Properties properties = new Properties();
        properties.put(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG, "300000");

        return KafkaSink.<E>builder()
                .setBootstrapServers(bs)
                .setRecordSerializer(KafkaRecordSerializationSchema.builder()
                        .setTopic(ot)
                        .setValueSerializationSchema(new CommonEntitySchema<>(clazz))
                        .build())
                .setKafkaProducerConfig(properties)
                .setDeliverGuarantee(DeliveryGuarantee.AT_LEAST_ONCE)
                .build();
    }

    /**
     * 构建kafka source
     * @param clazz 实体类class
     * @param <E> 实体类泛型
     * @return kafka source 对象
     */
    public <E> KafkaSource<E> buildSource(Class<E> clazz) throws ExecutionException, InterruptedException {
        final String kafkaConsumerGroup = parameterTool.getRequired(KafkaArgs.KAFKA_CONSUMER_GROUP.key);
        final String bootstrapServer = parameterTool.getRequired(KafkaArgs.BOOTSTRAP_SERVER.key);
        final String inputTopic = parameterTool.getRequired(KafkaArgs.INPUT_TOPIC.key);
        final boolean createTopic = parameterTool.has(KafkaArgs.CREATE_TOPIC.key);

        if (createTopic) {
            final int partition = parameterTool.getInt(KafkaArgs.CREATE_TOPIC_PARTITION.key, 1);
            final short replicationFactor = parameterTool.getShort(KafkaArgs.REPLICATION_FACTOR.key, (short) 1);
            createTopic(bootstrapServer, inputTopic, partition, replicationFactor);
        }

        return KafkaSource.<E>builder()
                .setGroupId(kafkaConsumerGroup)
                .setStartingOffsets(OffsetsInitializer.latest())
                .setBootstrapServers(bootstrapServer)
                .setTopics(inputTopic)
                .setValueOnlyDeserializer(new CommonEntitySchema<>(clazz))
                .build();
    }

    public enum KafkaArgs {
        /*
         * kafka 服务地址
         */
        BOOTSTRAP_SERVER("bs"),

        /*
         * kafka 消费者组
         */
        KAFKA_CONSUMER_GROUP("kcg"),

        /*
         * kafka 输入主题
         */
        INPUT_TOPIC("it"),

        /*
         * kafka 输出主题
         */
        OUTPUT_TOPIC("ot"),

        /*
         * 是否自动创建主题
         */
        CREATE_TOPIC("ct"),

        /*
         * 分区数
         */
        CREATE_TOPIC_PARTITION("pt"),

        /*
         * 副本数
         */
        REPLICATION_FACTOR("rf");

        private final String key;

        KafkaArgs(String key) {
            this.key = key;
        }
    }

    /**
     * 如果 TOPIC 不存在则创建该 TOPIC
     * @param bootstrapServer kafka broker 地址
     * @param topic 想要创建的 TOPIC
     * @param partitions 并行度
     * @param replicationFactor 副本数
     */
    public static void createTopic(String bootstrapServer,
                                   String topic,
                                   int partitions,
                                   int replicationFactor) throws ExecutionException, InterruptedException {
        Properties adminProperties = new Properties();
        adminProperties.setProperty(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer);
        try (AdminClient adminClient = AdminClient.create(adminProperties)) {
            if (!adminClient.listTopics().names().get().contains(topic)) {
                NewTopic newTopic = new NewTopic(topic, partitions, (short) replicationFactor);
                adminClient.createTopics(Collections.singletonList(newTopic)).all().get();
                log.info("created topic: {}", topic);
            }
        }
    }
}
